package com.edata.bigdata.nebula

import com.edata.bigdata.annotations.Edata_Saver
import com.edata.bigdata.basic.Saver
import com.vesoft.nebula.connector.connector.NebulaDataFrameWriter
import com.vesoft.nebula.connector.{WriteMode, WriteNebulaEdgeConfig, WriteNebulaVertexConfig}
import org.apache.spark.sql.{DataFrame, SparkSession}

@Edata_Saver(target = "NESAVER")
class NeSaver extends Saver with NeConnector {
  override var sourceType: String = _

  override var session: SparkSession = _


  private var vertexWriteConfigBuilder = WriteNebulaVertexConfig.builder()

  private var edgeWriteConfigBuilder = WriteNebulaEdgeConfig.builder()

  override def save(data: DataFrame, args: String*): Boolean = {
    val connConfig = createNebulaConnConfig()
    val schemaType = args(0)
    schemaType.toLowerCase match {
      case "vertex" => {
        val vertexWriteConfig = vertexWriteConfigBuilder
          .withUser(USER)
          .withPasswd(PASSWORD)
          .withWriteMode(WriteMode.INSERT).build()
        data.write.nebula(connConfig, vertexWriteConfig).writeVertices()
      }
      case "edge" => {
        val edgeWriteConfig = edgeWriteConfigBuilder
          .withUser(USER)
          .withPasswd(PASSWORD)
          .withWriteMode(WriteMode.INSERT).build()
        data.write.nebula(connConfig, edgeWriteConfig).writeEdges()
      }
    }
    return true
  }

  override def update(data: DataFrame, args: String*): Boolean = ???

  override def delete(data: DataFrame, args: String*): Boolean = ???

  override def saveOrUpdate(data: DataFrame, args: String*): Boolean = ???

  def setVertexWriterSpace(space_name: String): Unit = {
    vertexWriteConfigBuilder = vertexWriteConfigBuilder.withSpace(space_name)
  }

  def setVertexWriterTag(tag_name: String): Unit = {
    vertexWriteConfigBuilder = vertexWriteConfigBuilder.withTag(tag_name)
  }

  def setVertexWriterVidField(field_name: String): Unit = {
    vertexWriteConfigBuilder = vertexWriteConfigBuilder.withVidField(field_name)
  }

  def setVertexWriterVidAsProp(flag: Boolean): Unit = {
    vertexWriteConfigBuilder = vertexWriteConfigBuilder.withVidAsProp(flag)
  }

  def setVertexWriterBatch(num: Int): Unit = {
    vertexWriteConfigBuilder = vertexWriteConfigBuilder.withBatch(num)
  }

  def setVertexWriterDisableWriteLog(flag: Boolean): Unit = {
    vertexWriteConfigBuilder = vertexWriteConfigBuilder.withDisableWriteLog(flag)
  }

  def setEdgeWriterSpace(space_name: String): Unit = {
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withSpace(space_name)
  }

  def setEdgeWriterEdge(edge_name: String): Unit = {
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withEdge(edge_name)
  }

  def setEdgeWriterSrcIdField(field_name: String): Unit = {
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withSrcIdField(field_name)
  }

  def setEdgeWriterDstIdField(field_name: String): Unit = {
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withDstIdField(field_name)
  }

  def setEdgeWriterRankField(field_name: String): Unit = {
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withRankField(field_name)
  }

  def setEdgeWriterSrcAsProperty(flag:Boolean):Unit={
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withSrcAsProperty(flag)
  }

  def setEdgeWriterDstAsProperty(flag:Boolean):Unit={
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withDstAsProperty(flag)
  }

  def setEdgeWriterRankAsProperty(flag:Boolean):Unit={
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withRankAsProperty(flag)
  }

  def setEdgeWriterBatch(num:Int):Unit={
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withBatch(num)
  }

  def setEdgeWriterDisableWriteLog(flag:Boolean):Unit={
    edgeWriteConfigBuilder = edgeWriteConfigBuilder.withDisableWriteLog(flag)
  }




}
